import tensorflow as tf
from tensorflow import keras
import pandas as pd
import matplotlib.pyplot as plt
import time

if __name__ == '__main__':
    start = time.time()
    # 加载数据集
    fashion_mnist = keras.datasets.mnist
    (X_train_full, y_train_full), (X_test, Y_test) = fashion_mnist.load_data()
    # 查看训练集的形状和数据类型
    print(X_train_full.shape, X_train_full.dtype)
    # 比例缩放和像素强度降低到0-1,创建一个验证集
    X_valid, X_train = X_train_full[:5000] / 255.0, X_train_full[5000:] / 255.0
    Y_valid, Y_train = y_train_full[:5000], y_train_full[5000:]
    # 设置类名列表
    class_names = ["T_shirt/top", "Trouser", "Pullover", "Dress", "Coat", "Sandal", "shirt", "Sneaker"
        , "Bag", "Anlle boot"]
    # 查看第一幅图
    print(class_names[Y_train[0]])

    # 搭建网络模型
    model = tf.keras.models.Sequential()
    model.add(keras.layers.Flatten(input_shape=[28, 28]))
    model.add(keras.layers.Dense(300, activation="relu"))
    model.add(keras.layers.Dropout(0.4))
    model.add(keras.layers.Dense(100, activation="relu"))
    model.add(keras.layers.Dense(10, activation="softmax"))  # 输出十个概率分布，看属于哪一个
    model.compile(loss="sparse_categorical_crossentropy", optimizer="sgd", metrics=["accuracy"])
    history = model.fit(X_train, Y_train, epochs=5, validation_data=(X_valid, Y_valid), batch_size=200)
    # print(history)

    stop = time.time()
    print(stop - start)
    # 绘制history曲线
    pd.DataFrame(history.history).plot(figsize=(8, 5))
    plt.grid(True)
    plt.gca().set_ylim(0, 1)
    plt.show()

    # 在测试集上测试
    print(model.evaluate(X_test, Y_test))
    # 仅使用测试集的前三个例子
    X_new = X_test[:3]
    Y_proba = model.predict(X_new)
    print(Y_proba.round(2))
    y_pre = model.predict(X_new)
    print(y_pre)

